
import sqlite3
from datetime import datetime
from pathlib import Path
import hashlib
import secrets


DB_PATH = Path(__file__).parent / 'users.db'


class AuthDatabase:

    def __init__(self, db_path=None):
        self.db_path = db_path or DB_PATH

    def get_connection(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn



    def create_user(self, username, password, email, role='user'):
        conn = self.get_connection()
        cursor = conn.cursor()


        salt = secrets.token_hex(16)
        password_hash = self._hash_password(password, salt)


        default_xunfei_tokens = 10000
        default_qwen_tokens = 5000

        cursor.execute('''
            INSERT INTO users (
                username, password_hash, salt, email, role,
                xunfei_token_limit, qwen_token_limit,
                xunfei_tokens_used, qwen_tokens_used,
                is_active, created_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            username, password_hash, salt, email, role,
            default_xunfei_tokens, default_qwen_tokens,
            0, 0, 1, datetime.now().isoformat()
        ))

        user_id = cursor.lastrowid
        conn.commit()
        conn.close()

        return user_id

    def verify_user(self, username, password):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('SELECT * FROM users WHERE username = ?', (username,))
        user = cursor.fetchone()
        conn.close()

        if not user:
            return False, "用户不存在"

        if not user['is_active']:
            return False, "账号已被冻结"


        password_hash = self._hash_password(password, user['salt'])
        if password_hash != user['password_hash']:
            return False, "密码错误"


        self._update_last_login(user['id'])

        return True, {
            'id': user['id'],
            'username': user['username'],
            'email': user['email'],
            'role': user['role'],
            'xunfei_token_limit': user['xunfei_token_limit'],
            'qwen_token_limit': user['qwen_token_limit'],
            'xunfei_tokens_used': user['xunfei_tokens_used'],
            'qwen_tokens_used': user['qwen_tokens_used']
        }

    def get_user(self, user_id):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('SELECT * FROM users WHERE id = ?', (user_id,))
        user = cursor.fetchone()
        conn.close()

        if user:
            return dict(user)
        return None

    def get_all_users(self):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('''
            SELECT id, username, email, role, is_active,
                   xunfei_token_limit, qwen_token_limit,
                   xunfei_tokens_used, qwen_tokens_used,
                   created_at, last_login
            FROM users
            ORDER BY created_at DESC
        ''')

        users = [dict(row) for row in cursor.fetchall()]
        conn.close()

        return users

    def update_user_tokens(self, user_id, xunfei_limit=None, qwen_limit=None):
        conn = self.get_connection()
        cursor = conn.cursor()

        updates = []
        params = []

        if xunfei_limit is not None:
            updates.append('xunfei_token_limit = ?')
            params.append(xunfei_limit)

        if qwen_limit is not None:
            updates.append('qwen_token_limit = ?')
            params.append(qwen_limit)

        if updates:
            params.append(user_id)
            query = f"UPDATE users SET {', '.join(updates)} WHERE id = ?"
            cursor.execute(query, params)
            conn.commit()

        conn.close()

    def freeze_user(self, user_id):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('UPDATE users SET is_active = 0 WHERE id = ?', (user_id,))
        conn.commit()
        conn.close()

    def unfreeze_user(self, user_id):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('UPDATE users SET is_active = 1 WHERE id = ?', (user_id,))
        conn.commit()
        conn.close()

    def delete_user(self, user_id):
        conn = self.get_connection()
        cursor = conn.cursor()


        cursor.execute('DELETE FROM vocabulary WHERE user_id = ?', (user_id,))


        cursor.execute('DELETE FROM translation_history WHERE user_id = ?', (user_id,))


        cursor.execute('DELETE FROM users WHERE id = ?', (user_id,))

        conn.commit()
        conn.close()

    def check_username_exists(self, username):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('SELECT COUNT(*) as count FROM users WHERE username = ?', (username,))
        count = cursor.fetchone()['count']
        conn.close()

        return count > 0

    def check_email_exists(self, email):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('SELECT COUNT(*) as count FROM users WHERE email = ?', (email,))
        count = cursor.fetchone()['count']
        conn.close()

        return count > 0



    def record_token_usage(self, user_id, xunfei_tokens=0, qwen_tokens=0):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('''
            UPDATE users
            SET xunfei_tokens_used = xunfei_tokens_used + ?,
                qwen_tokens_used = qwen_tokens_used + ?
            WHERE id = ?
        ''', (xunfei_tokens, qwen_tokens, user_id))

        conn.commit()
        conn.close()

    def check_token_limit(self, user_id):
        user = self.get_user(user_id)

        if not user:
            return True, "用户不存在"

        if user['xunfei_tokens_used'] >= user['xunfei_token_limit']:
            return True, f"讯飞API额度已用完 ({user['xunfei_tokens_used']}/{user['xunfei_token_limit']})"

        if user['qwen_tokens_used'] >= user['qwen_token_limit']:
            return True, f"通义千问额度已用完 ({user['qwen_tokens_used']}/{user['qwen_token_limit']})"

        return False, "正常"

    def reset_user_tokens(self, user_id):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('''
            UPDATE users
            SET xunfei_tokens_used = 0, qwen_tokens_used = 0
            WHERE id = ?
        ''', (user_id,))

        conn.commit()
        conn.close()



    def add_translation_history(self, user_id, original_text, translated_text,
                               polished_text, direction, xunfei_usage, qwen_usage):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('''
            INSERT INTO translation_history (
                user_id, original_text, translated_text, polished_text,
                direction, xunfei_tokens, qwen_tokens, created_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        ''', (
            user_id, original_text, translated_text, polished_text,
            direction, xunfei_usage, qwen_usage, datetime.now().isoformat()
        ))

        conn.commit()
        conn.close()

    def get_user_history(self, user_id, limit=50):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('''
            SELECT * FROM translation_history
            WHERE user_id = ?
            ORDER BY created_at DESC
            LIMIT ?
        ''', (user_id, limit))

        history = [dict(row) for row in cursor.fetchall()]
        conn.close()

        return history



    def _hash_password(self, password, salt):
        return hashlib.sha256((password + salt).encode()).hexdigest()

    def _update_last_login(self, user_id):
        conn = self.get_connection()
        cursor = conn.cursor()

        cursor.execute('''
            UPDATE users SET last_login = ? WHERE id = ?
        ''', (datetime.now().isoformat(), user_id))

        conn.commit()
        conn.close()


def init_auth_db():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()


    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            salt TEXT NOT NULL,
            email TEXT UNIQUE NOT NULL,
            role TEXT DEFAULT 'user',
            xunfei_token_limit INTEGER DEFAULT 10000,
            qwen_token_limit INTEGER DEFAULT 5000,
            xunfei_tokens_used INTEGER DEFAULT 0,
            qwen_tokens_used INTEGER DEFAULT 0,
            is_active INTEGER DEFAULT 1,
            created_at TEXT NOT NULL,
            last_login TEXT
        )
    ''')


    cursor.execute('''
        CREATE TABLE IF NOT EXISTS translation_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            original_text TEXT NOT NULL,
            translated_text TEXT NOT NULL,
            polished_text TEXT,
            direction TEXT NOT NULL,
            xunfei_tokens INTEGER DEFAULT 0,
            qwen_tokens INTEGER DEFAULT 0,
            created_at TEXT NOT NULL,
            FOREIGN KEY (user_id) REFERENCES users(id)
        )
    ''')


    cursor.execute('''
        CREATE TABLE IF NOT EXISTS vocabulary_new (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            french TEXT NOT NULL,
            chinese TEXT NOT NULL,
            category TEXT DEFAULT '翻译结果',
            created_at TEXT NOT NULL,
            FOREIGN KEY (user_id) REFERENCES users(id)
        )
    ''')


    cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_id ON translation_history(user_id)')
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_vocab_user ON vocabulary_new(user_id)')
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_username ON users(username)')
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_email ON users(email)')

    conn.commit()


    auth_db = AuthDatabase(DB_PATH)
    if not auth_db.check_username_exists('Mike.Chen'):
        admin_id = auth_db.create_user(
            username='Mike.Chen',
            password='Chen@123',
            email='Chen@french-teacher.local',
            role='admin'
        )
        print(f"✓ 已创建默认管理员账号: Mike.Chen / Chen@123")

    conn.close()
    print(f"✓ 用户数据库已初始化: {DB_PATH}")


if __name__ == '__main__':

    print("="*60)
    print("🧪 用户数据库测试")
    print("="*60)

    init_auth_db()

    auth_db = AuthDatabase()


    if not auth_db.check_username_exists('testuser'):
        user_id = auth_db.create_user('testuser', 'password123', 'test@example.com')
        print(f"✓ 创建测试用户: ID {user_id}")


    success, result = auth_db.verify_user('admin', 'admin123')
    if success:
        print(f"✓ 登录成功: {result['username']} ({result['role']})")
    else:
        print(f"✗ 登录失败: {result}")


    users = auth_db.get_all_users()
    print(f"\n用户列表 ({len(users)} 个):")
    for user in users:
        print(f"  - {user['username']} ({user['role']}) - Token: {user['xunfei_tokens_used']}/{user['xunfei_token_limit']}")

    print("\n✅ 测试完成")
